# 本程序实现 对刚度矩阵 和 载荷向量的 读写操作
# 以P1元作为操作示例
# 达到类似何晓明 P_T T_b矩阵类似的效果
# HeXiaoMing书 3.7.1 Example1
# -Laplace(u) = f in \Omega
# u = g on \partial\Omega
import importlib.util
import tools
import norm

if importlib.util.find_spec("petsc4py") is not None:
    import dolfinx
    if not dolfinx.has_petsc:
        print("This demo requires DOLFINx to be compiled with PETSc enabled.")
        exit(0)
    from petsc4py.PETSc import ScalarType  # type: ignore
else:
    print("This demo requires petsc4py.")
    exit(0)

from mpi4py import MPI
from basix.ufl import element
import numpy as np
from petsc4py import PETSc
import ufl
from dolfinx import fem, io, mesh, plot, default_real_type
from dolfinx.fem.petsc import LinearProblem, assemble_matrix, assemble_vector, create_matrix
from ufl import ds, dx, grad, inner
from dolfinx.fem.petsc import set_bc  # 设置边界条件新api


def getStiffnessMatrix(A: PETSc.Mat):
    """
    获取刚度矩阵的值(np.array)
    """
    # 需要导入
    # from petsc4py import PETSc
    m, n = A.getSize()
    matrix_data = np.zeros((m, n))

    # 获取矩阵的所有非零元素
    row, col, val = A.getValuesCSR()

    # 将CSR格式的数据转换为密集矩阵
    for i in range(m):
        for j in range(row[i], row[i + 1]):
            matrix_data[i, col[j]] = val[j]

    return matrix_data


# cell_n = [1, 1]
# cell_n = [4, 4]
cell_n = [8, 8]
# cell_n = [16, 16]
# +
msh = mesh.create_rectangle(
    comm=MPI.COMM_WORLD,
    points=((-1, -1), (1.0, 1.0)),
    n=cell_n,
    cell_type=mesh.CellType.triangle,
)
u_p1 = element("Lagrange", msh.basix_cell(),
               degree=1, dtype=default_real_type)
V = fem.functionspace(msh, u_p1)


facets = mesh.locate_entities_boundary(
    msh,
    dim=(msh.topology.dim - 1),
    marker=lambda x: np.isclose(x[0], -1) | np.isclose(
        x[0], 1.0) | np.isclose(x[1], -1) | np.isclose(x[1], 1.0)
)
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
x = ufl.SpatialCoordinate(msh)


def u_true(x):
    return np.exp(x[0] + x[1])


def u_true_(x, y):
    return np.exp(x + y)


def partial_derivative_x(x, y):
    """
    计算函数 u 对 x 的偏导数的数值
    :param x: 自变量 x 的值
    :param y: 自变量 y 的值
    :return: 偏导数 ∂u/∂x 的数值结果
    """
    return np.exp(x + y)


def partial_derivative_y(x, y):
    """
    计算函数 u 对 y 的偏导数的数值
    :param x: 自变量 x 的值
    :param y: 自变量 y 的值
    :return: 偏导数 ∂u/∂y 的数值结果
    """
    return np.exp(x + y)

def grad_u(x, y):
    return np.array([partial_derivative_x(x, y), partial_derivative_y(x, y)])


def g_partial(x, y):
    u_dx = partial_derivative_x(x, y)
    u_dy = partial_derivative_y(x, y)
    if np.all(np.isclose(x, -1)):
        res = -u_dx
    elif np.all(np.isclose(x, 1)):
        res = u_dx
    elif np.all(np.isclose(y, -1)):
        res = -u_dy
    elif np.all(np.isclose(y, 1)):
        res = u_dy
    else:
        raise ValueError("边界配置错误")
    return res


def f_np(x, y):            # 指数部分
    return -2 * np.exp(x + y)


f = ufl.exp(x[0] + x[1])

a = inner(grad(u), grad(v)) * dx  # <class 'ufl.form.Form'>
a_form = fem.form(a)  # < class 'dolfinx.fem.forms.Form' >
L = inner(f, v) * dx
L_form = fem.form(L)

u_err = dolfinx.fem.Function(V)
u_exact = dolfinx.fem.Function(V)
u_exact.interpolate(u_true)

dofs = fem.locate_dofs_topological(V=V, entity_dim=1, entities=facets)
bc = fem.dirichletbc(u_exact, dofs=dofs)


A = assemble_matrix(a_form, [bc])   # 刚度矩阵
A.assemble()
b = assemble_vector(L_form)  # 载荷向量 需要手动处理自由度

fem.apply_lifting(b, [a_form], [[bc]])  # 调整载荷项b，考虑边界条件影响
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
set_bc(b, [bc])

A_mat = getStiffnessMatrix(A)  # ndarray


# He XiaoMing程序思路
G = tools.GaussIntegrate(order=9)  # 处理高斯积分计算
dof_map = V.dofmap  # 自由度映射
cells = 2 * cell_n[0] ** 2

# 手动组装
A_hxm = np.zeros_like(A_mat)
A1 = np.zeros_like(A_mat)
A2 = np.zeros_like(A_mat)
b_hxm = np.zeros_like(b.array)


# dolfinx返回的第三维是z空间 二维问题直接取0
# 虚拟区域法论文 使用的是重心坐标 所以积分公式不能直接转化!!!

for i in range(cells):
    dof_idx = dof_map.cell_dofs(i)  # i单元对应的顶点全局索引
    nodes = V.tabulate_dof_coordinates()[dof_idx]  # 顶点的物理坐标
    
    bdt = tools.BasisDerivativeTransformer(nodes[:, :-1])
    J = bdt.J
    J_inv = np.linalg.inv(J)  # 计算Jacobian矩阵的逆
    det_J = bdt.detJ

    basis_func_res = u_p1.tabulate(1, G.points)
    gs_basis_value = basis_func_res[0]          # 函数值
    diff_gs_basis_value_xi = basis_func_res[1]  # ξ 导数
    diff_gs_basis_value_eta = basis_func_res[2]  # η 导数
    gs_points_phy = G.to_physical_gauss_points(*nodes[:, :-1])  # 实际单元高斯点

    # 计算刚度矩阵
    for alpha in range(len(dof_idx)):  # 遍历P1元基函数 3个; P2则为6个
        for beta in range(len(dof_idx)):
            # 刚度矩阵 索引
            # 存储物理单元alpha和beta索引下 的导函数值
            
            local_diff_alpha_x = bdt.local_dx1(
                diff_gs_basis_value_xi[:, alpha], diff_gs_basis_value_eta[:, alpha])
            local_diff_alpha_y = bdt.local_dy1(
                diff_gs_basis_value_xi[:, alpha], diff_gs_basis_value_eta[:, alpha])
            local_diff_beta_x = bdt.local_dx1(
                diff_gs_basis_value_xi[:, beta], diff_gs_basis_value_eta[:, beta])
            local_diff_beta_y = bdt.local_dy1(
                diff_gs_basis_value_xi[:, beta], diff_gs_basis_value_eta[:, beta])
            # 计算刚度矩阵
            rr = local_diff_alpha_x * local_diff_beta_x + local_diff_alpha_y * local_diff_beta_y
            A_hxm[dof_idx[beta], dof_idx[alpha]] += np.dot(rr, G.weights) * det_J

    for beta in range(len(dof_idx)):  # 遍历P1元基函数 3个; P2则为6个
        b_hxm[dof_idx[beta]] += np.dot(f_np(gs_points_phy[:, 0], gs_points_phy[:, 1])
                                       * gs_basis_value[:, beta], G.weights) * det_J

# Neumann边界条件
# 边界只有 y == -1
num_basis_u = u_p1.degree * 3
G1d = tools.GaussIntegrate1D(order=9)
boundary_dict = tools.search_boundary_edge_cell(
    msh, counterclock=True)
for bd_edge, edge_info in boundary_dict.items():
    p1, p2 = edge_info["coordinates"]
    if np.isclose(p1[1], -1) and np.isclose(p2[1], -1):  
        cell_idx = edge_info["cell_index"]
        u_idx = dof_map.cell_dofs(cell_idx)
        nodes = V.tabulate_dof_coordinates()[u_idx]  # 顶点的物理坐标
        bdt = tools.BasisDerivativeTransformer(nodes[:, :-1])
        det_J = bdt.detJ
        # 对边界单元进行遍历coordinates
        a, b = np.sort([p1[0], p2[0]])
        
        
        gs_points_phy = np.zeros((2, G1d.points.size))
        gs_points_phy[0] = G1d.get_nodes(p1[0], p2[0])
        gs_points_phy[1] = -1
        ref_points = bdt.local_to_ref(*gs_points_phy)
        f_v = g_partial(*gs_points_phy)
        
        for beta in range(num_basis_u):
            rr = f_v * u_p1.tabulate(0, ref_points.T)[0][:, beta]
            b_hxm[u_idx[beta]] += G1d.integrate(a, b, rr)
 

for dof in dofs:
    coords = V.tabulate_dof_coordinates()[dof][:-1]  # 获取前两维坐标
    # 检查是否在 y=1、x=1 或 x=-1 边界上
    is_boundary = (np.isclose(coords[1], 1) or
                   np.isclose(coords[0], 1) or
                   np.isclose(coords[0], -1))

    if is_boundary:
        # 施加 Dirichlet 边界条件
        A_hxm[dof, :] = 0
        A_hxm[dof, dof] = 1
        b_hxm[dof] = u_exact.x.array[dof]

print(f"A_hxm形状: {A_hxm.shape}")  # 应输出 (9, 9)
print(f"b_hxm形状: {b_hxm.shape}")  # 应输出 (9,)
u_manual = tools.sparse_solve(A_hxm, b_hxm)



# 调用norm.py 计算误差
u_L2_err = norm.get_L2_err(
    msh=msh,
    dof_map=dof_map,
    num_cells=cells,
    u_true_func=u_true,
    u_num=u_manual,
    func_space=V,
    order=1
)
u_max_err = norm.get_max_node_err(
    u_true_func=u_true,
    u_num=u_manual,
    func_space=V
)
u_h1_err = norm.get_H1_err(
    msh=msh,
    dof_map=dof_map,
    num_cells=cells,
    u_true_func_diff_ls=[partial_derivative_x, partial_derivative_y],
    u_num=u_manual,
    func_space=V,
    order=1
)
print(f"maxNode Err = {u_max_err}")
print(f"L2 err = {u_L2_err}")
print(f"H1 err = {u_h1_err}")